# coding=utf-8
import re
from ml.naivebeyes.beyes import *
from numpy import array


def parse_text(email_content):
    token = re.split(r'\W*', email_content)
    return [x.lower() for x in token if len(x) > 2]


doc_list = []
class_list = []
words = []

for i in range(1, 26):
    word_list = parse_text(open('email/spam/%d.txt' % i).read())
    doc_list.append(word_list)
    words.extend(word_list)
    class_list.append(1)

    word_list = parse_text(open('email/ham/%d.txt' % i).read())
    doc_list.append(word_list)
    words.extend(word_list)
    class_list.append(0)

vocab = create_vocab_list(doc_list)
train_set = range(50)
test_set = []

for i in range(10):
    index = int(random.uniform(0, len(train_set)))
    test_set.append(index)
    del train_set[index]

train_mat = []
train_class = []

for i in train_set:
    train_mat.append(words2vec(vocab, doc_list[i]))
    train_class.append(class_list[i])

p0, p1, ps = train_nb0(array(train_mat), array(train_class))

error_count = 0
for i in test_set:
    word_vec = words2vec(vocab, doc_list[i])
    if classify_nb(array(word_vec), p0, p1, ps) != class_list[i]:
        error_count += 1

print("the error rate is %f" % (float(error_count) / len(test_set)))
